import pickle
import threading
from typing import Optional, Tuple

import numpy as np
import zmq
from typing import Optional, Protocol, Tuple
import json

DEFAULT_CAMERA_PORT = 5000

class ZMQClientCamera:
    """A class representing a ZMQ client for a leader robot."""

    def __init__(self, port: int = DEFAULT_CAMERA_PORT, host: str = "127.0.0.1"):
        self._context = zmq.Context()
        self._socket = self._context.socket(zmq.SUB)
        self._socket.connect(f"tcp://{host}:{port}")
        self._socket.setsockopt(zmq.SUBSCRIBE, b"")
        self._socket.setsockopt(zmq.RCVHWM, 2)
        self._socket.setsockopt(zmq.CONFLATE, 1)

    def read(self) -> Tuple[np.ndarray]:
        """Get the current state of the leader robot.

        Returns:
            T: The current state of the leader robot.
        """
        # message = self._socket.recv()
        # print(message)
        # image = pickle.loads(message)
        # return image

        msg = self._socket.recv_multipart()
        metadata = json.loads(msg[0].decode('utf-8'))
        buffer = msg[1]

        # 重建numpy数组
        image = np.frombuffer(buffer, dtype=np.dtype(metadata['dtype']))
        image = image.reshape(metadata['shape'])

        return image


def main():
    ZMQClientCamera()


if __name__ == '__main__':
    main()
